#include <stdlib.h>
#include <stdio.h>
#include <string>
#include <fstream>
#include <sstream>
#include <iostream>
#include <iomanip>
#include <sys/stat.h>
#include <vector>
#include <signal.h>
#include "sockets/ClientSocket.h"
#include "sockets/SocketException.h"
#include "Constants.h"

using namespace std;
string PUBKEYFILE = "cloudkey";

bool checkfile(string);
string int2str(int);
int str2int(string);
void cleanrun(string,string,string);
void cleanup();
void parsehostfile(string);
void writeresults(string);
void setupmptrace(ClientSocket* sock);

//For capturing ctrl+c kill, to force close properly
void sigproc(int sig);

int main(int argv, char** argc){

  int newport;
  string newhost,t;
  string data,pkfile,mpicom;
  string hostfile,resultfile,codefile;
  int numnodes;
  string command,checkproc;
  vector<string> pubkeyfiles; //for master node
  ClientSocket stable_socket;
  ifstream pubkey;

  //Capture the kill signal to close properly
  signal(SIGINT,&sigproc);

  //Client properties.
  int masterflag=0;

  try
    {
      //The goal in this block is to get the newport and newhost variables set.
      ClientSocket init_socket ( DEFAULT_HOST, DEFAULT_PORT );
      string data;
      string reply;
      string temp;
      
      try
	{
	  init_socket << (string)"Test message.";
	  init_socket >> data;
	}
      catch ( SocketException& ) {}
      int loc1,loc2;
      loc1=data.find('-',0);
      reply=data.substr(0,loc1);
      loc2=data.find('-',loc1+1);
      newhost=data.substr(loc1+1,loc2-loc1-1);
      newport=atoi((data.substr(loc2+1)).c_str());
      //Debug
      //printf("reply:%s\nnewhost:%s\nnewport:%d\n",reply.c_str(),newhost.c_str(),newport);
      
      if(reply.compare("Message recieved.") != 0){
	//kill the connection, exit with error
	init_socket.close();
	printf("Wrong host response recieved (%s), possibly connecting to wrong host.\n",reply.c_str());
	return -1;
      } 
      
      //Close the initial socket and open a new one on the host/port specified.
      init_socket.close();
      
    }
  catch ( SocketException& e )
    {
      cout << "Client exception was caught:" << e.description() << "\n";
      cout << "Server may not exist or connection to server may not be possible.\n";
      return -1;
    }



  //If you've made it this far, you have a new port & host to connect to.
  //Below block sets up the stable_socket
  try
    {
      printf("Trying to make a stable connection...\n");
      int count=0;
      while(count<12){ //wait up to 12*5 seconds before giving up on connecting
	sleep(5);
	try{
	  printf("Trying to connect to %s on port %d\n",newhost.c_str(),newport);
	  stable_socket.init( newhost, newport );
	  printf("stable socket created\n");
	  break;
	} catch (SocketException&) {}
	count++;
      }
      if(count==12){
	printf("Error: Unable to connect to stable host - host connection timeout!\n");
	return -1;
      } 
      
      //Sending confirmation message
      try
	{
	  stable_socket << (string)"You there?";
	  stable_socket >> data;
	}
      catch ( SocketException& ) {}
      
      if(data.compare("Yessir.") != 0){
	//kill the connection, exit with error
	stable_socket.close();
	printf("Wrong stable host response recieved, possibly connecting to wrong host.\n");
	return -1;
      } else {
	//Stable Connection created
	printf("Stable connection created!\n");
	
	//Generate Public key, pass to server
	PUBKEYFILE=PUBKEYFILE+int2str(newport)+".pub";
	cout << "Generating temporary public keyfile...\n";
	t=GENPASS+" "+int2str(newport);
	system(t.c_str());
	sleep(1);

	//Transfer the key
	stable_socket.sendfile(PUBKEYFILE);

	//Cleanup
	t="rm "+PUBKEYFILE;
	system(t.c_str());

	//Send the server the client username.
	system("whoami > username");
	ifstream username;
	username.open("username");
	username>>t;
	stable_socket<<t;
	t="rm username";
	system(t.c_str());
      }
     
    }
  catch ( SocketException& e )
    {
      cout << "Client (stable_connection) exception was caught:" << e.description() << "\n";
      raise(SIGINT);
    }

  //Client connection up, loop here waiting for more instructions etc.
  //start mpd

  //Move below to a thread, make stable_socket global (or pass as arguement)
  try{
    while(true){
      //sleep(5); //MASTER SLEEP STATEMENT
      stable_socket>>data;
      printf("Switch!\n");
      switch(str2int(data)) {  //Compare server command

	//=========================
      case TICK : //Heartbeat
	printf("Heartbeat::TickTock\n");
	if(masterflag){
	  //I am the master node and the job is running.
	  system("rm psfile");
	  system("ps aux > psfile");
	  system("rm psresults");
	  checkproc="grep "+mpicom+" psfile > psresults";
	  system(checkproc.c_str());
	  //Fake calculation delay
	  //sleep(10);
	  //Write a fake resultfile;
	  //writeresults(resultfile);
	  if(checkfile((string)"psresults")){
	    //Job is still running.
	    stable_socket<<int2str(RUNS);
	  } else {
	    //Job is complete.
	    masterflag=0;
	    if(checkfile(resultfile)){
	      printf("SUCCESS! Sending back results.\n");
	      stable_socket<<int2str(DONE);
	      system(command.c_str());
	      stable_socket.sendfile(resultfile);
	    } else {
	      stable_socket<<int2str(TOCK);
	    }
	    cleanrun(resultfile,codefile,hostfile);
	  }
	} else {
	  //IDLE
	  stable_socket<<int2str(TOCK);
	}
	break;
	
	//=========================
      case NBSY :
	cout << "Receiving job from server\n";
	stable_socket<<int2str(ACKN);	
	
	// Transfering codefile  
	cout << "Receiving the codefile.\n";
	codefile=stable_socket.recvfile();

	// Transfering hostfile  
	cout << "Receiving the hostfile.\n";
	hostfile=stable_socket.recvfile();

	printf("getting numnodes\n");

	// Transfering the number of nodes
	stable_socket>>data;
	numnodes=str2int(data);
	stable_socket<<int2str(ACKN);

	printf("getting resultfile name\n");

	//Transfering resultfile name
	stable_socket >> resultfile;
	
	//Parse the host file
	parsehostfile(hostfile);

	printf("setting up mpd on master.\n");

	//Call mpdsetup.sh script
	command=MPDSETUP;
	system(command.c_str());

	printf("calling mpd on master.\n");
	command="mpd";
	system(command.c_str());

	printf("calling mpdtrace on master.\n");
	setupmptrace(&stable_socket);

	//Wait for workers to be ready.
	stable_socket >> data;

	sleep(1);
	
	printf("running mpi\n");

	//Start the code.
	mpicom="mpdmpexec -n "+int2str(numnodes)+" ./"+codefile;
	system(mpicom.c_str());
	masterflag=1;
	break;
	
      case WRKR :
	//Recieve the public key
	cout << "Recieving master public key file.\n";
	data=stable_socket.recvfile();
	//Call addkeys.sh script
	t=ADDKEYS+" "+data;
	system(t.c_str());
	//Clean up the key now that its added to Auth list.
	t="rm "+data;
	system(t.c_str());
	
	//Call mpdsetup.sh script
	command=MPDSETUP;
	system(command.c_str());

	stable_socket >> data;
	stable_socket << int2str(ACKN);
	t=data;//special number
	stable_socket >> data;
	command="mpd -h "+data+" -p "+t;
	cout << "MPD COMMAND ON SLAVE: " << command;
	system(command.c_str());
	break;
	
      default :
	printf("Error in switch, should never get here!\n");
	
      }
    }
  }
  catch(SocketException& e ) {
    cout << "Connection to Server lost! Disconnecting and stopping run!  Error: " << e.description() << "\n";
    //End job here
    raise(SIGINT);
  }
 
  return 0;
}

//Returns true if file exists
bool checkfile(string lockfilename){
  struct stat lfinfo;
  int st;
  st=stat(lockfilename.c_str(),&lfinfo);
  if(st==0)
    return true;
  return false;
}

string int2str(int i){
  stringstream out;
  out << i;
  return out.str();
}

int str2int(string s){
  return atoi(s.c_str());
}

void cleanup(){
  //Remove the keyfiles
  //system("rm -rf keys/*");
  //system("rm -f psfile psresults");
  //string com="rm -f"+PUBKEYFILE;
  //system(com.c_str());
  //Run wipepass to change authorized keys back
  system(WIPEPASS.c_str());
}

void cleanrun(string results,string codefile, string hostfile){
  //delete keys and results etc
  string command;
  command="rm -f "+results+" "+codefile+" "+hostfile;
  system(command.c_str());
}

void sigproc(int sig){ 	
  if(sig==SIGINT){
    cout << "\nCleaning up...\n";
    cleanup();
    cout << "Closing.\n";
    exit(0);
  }
}

void parsehostfile(string hostfile){
  ifstream hfile;
  hfile.open(hostfile.c_str());
  char buff[128];
  string line,ip,uname;
  string command;
  int loc;
  vector<string> ipuns;

  hfile.getline(buff,128);
  while(!hfile.eof()){
    //Parse out the uname and IP address
    line=buff;
    loc=line.find("@",0);
    uname=line.substr(0,loc);
    ip=line.substr(loc+1,line.size()-loc-1);
    cout << uname << " " << ip << endl;
    command=ADDHOST+" "+uname+" "+ip;
    system(command.c_str());
    ipuns.push_back(ip);
    hfile.getline(buff,128);
  }
  hfile.close();

  ofstream nhfile;
  nhfile.open(hostfile.c_str());
  for(int i=0;i<(int)ipuns.size();i++){
    nhfile<< (string)ipuns[i];
    nhfile<< "\n";
  }
  nhfile.close();
}

void writeresults(string resfile){
  ofstream results;
  results.open(resfile.c_str());
  results << "If you want your results, depost 20$ in to paypal account 19284291.\n";
  results << "\n\n";
  results << "If your results are not displayed properly please contact stashu@gmail.com for assistance.\n";
  results.close();
}

void setupmptrace(ClientSocket* sock){
  //string command="mpdtrace > mtfile";
  char buffer[128];
  ifstream mt;
  int loc1,loc2;
  mt.open("mtfile");
  mt.getline(buffer,128);
  string line=buffer;
  loc1=line.find("_");
  loc2=line.find(":",loc1+1);
  string num=line.substr(loc1+1,loc2-loc1-1);
  *sock << num;
  //system("rm mtfile");
}
